This tutorial is part of multiphenicsx.

Tutorial 06, case 5: Stokes problem with distributed control

In this tutorial we solve the optimal control problem

$$\min J(y, u) = \frac{1}{2} \int_{\Omega} |v - v_d|^2 dx + \frac{\alpha}{2} \int_{\Omega} |u|^2 dx$$ s.t. $$\begin{cases} - \Delta v + \nabla p = f + u & \text{in } \Omega\\ \text{div} v = 0 & \text{in } \Omega\\ v = 0 & \text{on } \partial\Omega \end{cases}$$

where $$\begin{align*} & \Omega & \text{unit square}\\ & u \in [L^2(\Omega)]^2 & \text{control variable}\\ & v \in [H^1_0(\Omega)]^2 & \text{state velocity variable}\\ & p \in L^2(\Omega) & \text{state pressure variable}\\ & \alpha > 0 & \text{penalization parameter}\\ & v_d & \text{desired state}\\ & f & \text{forcing term} \end{align*}$$ using an adjoint formulation solved by a one shot approach

import dolfinx.fem
import dolfinx.fem.petsc
import dolfinx.io
import dolfinx.mesh
import mpi4py.MPI
import numpy as np
import numpy.typing
import petsc4py.PETSc
import sympy
import ufl
import viskex

import multiphenicsx.fem
import multiphenicsx.fem.petsc

Mesh

mesh = dolfinx.mesh.create_unit_square(mpi4py.MPI.COMM_WORLD, 32, 32)

def bottom(x: np.typing.NDArray[np.float64]) -> np.typing.NDArray[np.bool_]:
    """Condition that defines the bottom boundary."""
    return abs(x[1] - 0.) < np.finfo(float).eps  # type: ignore[no-any-return]


def left(x: np.typing.NDArray[np.float64]) -> np.typing.NDArray[np.bool_]:
    """Condition that defines the left boundary."""
    return abs(x[0] - 0.) < np.finfo(float).eps  # type: ignore[no-any-return]


def top(x: np.typing.NDArray[np.float64]) -> np.typing.NDArray[np.bool_]:
    """Condition that defines the top boundary."""
    return abs(x[1] - 1.) < np.finfo(float).eps  # type: ignore[no-any-return]


def right(x: np.typing.NDArray[np.float64]) -> np.typing.NDArray[np.bool_]:
    """Condition that defines the right boundary."""
    return abs(x[0] - 1.) < np.finfo(float).eps  # type: ignore[no-any-return]


boundaries_entities = dict()
boundaries_values = dict()
for (boundary, boundary_id) in zip((bottom, left, top, right), (1, 2, 3, 4)):
    boundaries_entities[boundary_id] = dolfinx.mesh.locate_entities_boundary(
        mesh, mesh.topology.dim - 1, boundary)
    boundaries_values[boundary_id] = np.full(
        boundaries_entities[boundary_id].shape, boundary_id, dtype=np.int32)
boundaries_entities_unsorted = np.hstack(list(boundaries_entities.values()))
boundaries_values_unsorted = np.hstack(list(boundaries_values.values()))
boundaries_entities_argsort = np.argsort(boundaries_entities_unsorted)
boundaries_entities_sorted = boundaries_entities_unsorted[boundaries_entities_argsort]
boundaries_values_sorted = boundaries_values_unsorted[boundaries_entities_argsort]
boundaries = dolfinx.mesh.meshtags(
    mesh, mesh.topology.dim - 1,
    boundaries_entities_sorted, boundaries_values_sorted)
boundaries.name = "boundaries"

boundaries_1234 = boundaries.indices[np.isin(boundaries.values, (1, 2, 3, 4))]

viskex.dolfinx.plot_mesh(mesh)

viskex.dolfinx.plot_mesh_entities(mesh, mesh.topology.dim - 1, "boundaries", boundaries_1234)

Function spaces

Y_velocity = dolfinx.fem.functionspace(mesh, ("Lagrange", 2, (mesh.geometry.dim, )))
Y_pressure = dolfinx.fem.functionspace(mesh, ("Lagrange", 1))
U = dolfinx.fem.functionspace(mesh, ("Lagrange", 2, (mesh.geometry.dim, )))
Q_velocity = Y_velocity.clone()
Q_pressure = Y_pressure.clone()

Trial and test functions

(v, p) = (ufl.TrialFunction(Y_velocity), ufl.TrialFunction(Y_pressure))
(w, q) = (ufl.TestFunction(Y_velocity), ufl.TestFunction(Y_pressure))
u = ufl.TrialFunction(U)
r = ufl.TestFunction(U)
(z, b) = (ufl.TrialFunction(Q_velocity), ufl.TrialFunction(Q_pressure))
(s, d) = (ufl.TestFunction(Q_velocity), ufl.TestFunction(Q_pressure))

Problem data

alpha = 1.e-5
epsilon = 1.e-5
x, y = sympy.symbols("x[0], x[1]")
psi_d = 10 * (1 - sympy.cos(0.8 * np.pi * x)) * (1 - sympy.cos(0.8 * np.pi * y)) * (1 - x)**2 * (1 - y)**2
v_d_x = sympy.lambdify([x, y], psi_d.diff(y, 1))
v_d_y = sympy.lambdify([x, y], - psi_d.diff(x, 1))
v_d = dolfinx.fem.Function(Y_velocity)
v_d.interpolate(lambda x: np.stack((v_d_x(x[0], x[1]), v_d_y(x[0], x[1])), axis=0))
ff = dolfinx.fem.Constant(mesh, tuple(petsc4py.PETSc.ScalarType(0) for _ in range(2)))
bc0 = np.zeros((2, ), dtype=petsc4py.PETSc.ScalarType)

Optimality conditions

a = [[ufl.inner(v, w) * ufl.dx, None, None, ufl.inner(ufl.grad(z), ufl.grad(w)) * ufl.dx,
      - ufl.inner(b, ufl.div(w)) * ufl.dx],
     [None, None, None, - ufl.inner(ufl.div(z), q) * ufl.dx, epsilon * ufl.inner(b, q) * ufl.dx],
     [None, None, alpha * ufl.inner(u, r) * ufl.dx, - ufl.inner(z, r) * ufl.dx, None],
     [ufl.inner(ufl.grad(v), ufl.grad(s)) * ufl.dx, - ufl.inner(p, ufl.div(s)) * ufl.dx,
      - ufl.inner(u, s) * ufl.dx, None, None],
     [- ufl.inner(ufl.div(v), d) * ufl.dx, epsilon * ufl.inner(p, d) * ufl.dx, None, None, None]]
f = [ufl.inner(v_d, w) * ufl.dx,
     None,
     None,
     ufl.inner(ff, s) * ufl.dx,
     None]
a[3][3] = dolfinx.fem.Constant(mesh, petsc4py.PETSc.ScalarType(0)) * ufl.inner(z, s) * ufl.dx
f[1] = ufl.inner(dolfinx.fem.Constant(mesh, petsc4py.PETSc.ScalarType(0)), q) * ufl.dx
f[2] = ufl.inner(dolfinx.fem.Constant(mesh, tuple(petsc4py.PETSc.ScalarType(0) for _ in range(2))), r) * ufl.dx
f[4] = ufl.inner(dolfinx.fem.Constant(mesh, petsc4py.PETSc.ScalarType(0)), d) * ufl.dx
a_cpp = dolfinx.fem.form(a)
f_cpp = dolfinx.fem.form(f)
bdofs_Y_velocity_1234 = dolfinx.fem.locate_dofs_topological(
    Y_velocity, mesh.topology.dim - 1, boundaries_1234)
bdofs_Q_velocity_1234 = dolfinx.fem.locate_dofs_topological(
    Q_velocity, mesh.topology.dim - 1, boundaries_1234)
bc = [dolfinx.fem.dirichletbc(bc0, bdofs_Y_velocity_1234, Y_velocity),
      dolfinx.fem.dirichletbc(bc0, bdofs_Q_velocity_1234, Q_velocity)]

Solution

(v, p) = (dolfinx.fem.Function(Y_velocity), dolfinx.fem.Function(Y_pressure))
u = dolfinx.fem.Function(U)
(z, b) = (dolfinx.fem.Function(Q_velocity), dolfinx.fem.Function(Q_pressure))

Cost functional

J = 0.5 * ufl.inner(v - v_d, v - v_d) * ufl.dx + 0.5 * alpha * ufl.inner(u, u) * ufl.dx
J_cpp = dolfinx.fem.form(J)

Uncontrolled functional value

# Extract state forms from the optimality conditions
a_state = [[ufl.replace(a[i][j], {s: w, d: q}) if a[i][j] is not None else None
            for j in (0, 1)] for i in (3, 4)]
f_state = [ufl.replace(f[i], {s: w, d: q}) for i in (3, 4)]
a_state_cpp = dolfinx.fem.form(a_state)
f_state_cpp = dolfinx.fem.form(f_state)
bc_state = [bc[0]]

# Assemble the block linear system for the state
A_state = dolfinx.fem.petsc.assemble_matrix_block(a_state_cpp, bcs=bc_state)
A_state.assemble()
F_state = dolfinx.fem.petsc.assemble_vector_block(f_state_cpp, a_state_cpp, bcs=bc_state)

# Solve
vp = dolfinx.fem.petsc.create_vector_block([f_cpp[j] for j in (0, 1)])
ksp = petsc4py.PETSc.KSP()
ksp.create(mesh.comm)
ksp.setOperators(A_state)
ksp.setType("preonly")
ksp.getPC().setType("lu")
ksp.getPC().setFactorSolverType("mumps")
ksp.setFromOptions()
ksp.solve(F_state, vp)
vp.ghostUpdate(addv=petsc4py.PETSc.InsertMode.INSERT, mode=petsc4py.PETSc.ScatterMode.FORWARD)
ksp.destroy()

<petsc4py.PETSc.KSP at 0x7efd2d624680>

# Split the block solution in components
with multiphenicsx.fem.petsc.BlockVecSubVectorWrapper(vp, [c.function_space.dofmap for c in (v, p)]) as vp_wrapper:
    for vp_wrapper_local, component in zip(vp_wrapper, (v, p)):
        with component.vector.localForm() as component_local:
            component_local[:] = vp_wrapper_local

J_uncontrolled = mesh.comm.allreduce(dolfinx.fem.assemble_scalar(J_cpp), op=mpi4py.MPI.SUM)
print("Uncontrolled J =", J_uncontrolled)
assert np.isclose(J_uncontrolled, 0.1784536)

Uncontrolled J = 0.17845361493420545

viskex.dolfinx.plot_vector_field(v, "uncontrolled state velocity", glyph_factor=1e-1)

viskex.dolfinx.plot_scalar_field(p, "uncontrolled state pressure")

Optimal control

# Assemble the block linear system for the optimality conditions
A = dolfinx.fem.petsc.assemble_matrix_block(a_cpp, bcs=bc)
A.assemble()
F = dolfinx.fem.petsc.assemble_vector_block(f_cpp, a_cpp, bcs=bc)

# Solve
vpuzb = dolfinx.fem.petsc.create_vector_block(f_cpp)
ksp = petsc4py.PETSc.KSP()
ksp.create(mesh.comm)
ksp.setOperators(A)
ksp.setType("preonly")
ksp.getPC().setType("lu")
ksp.getPC().setFactorSolverType("mumps")
ksp.setFromOptions()
ksp.solve(F, vpuzb)
vpuzb.ghostUpdate(addv=petsc4py.PETSc.InsertMode.INSERT, mode=petsc4py.PETSc.ScatterMode.FORWARD)
ksp.destroy()

<petsc4py.PETSc.KSP at 0x7efd2d50e890>

# Split the block solution in components
with multiphenicsx.fem.petsc.BlockVecSubVectorWrapper(
        vpuzb, [c.function_space.dofmap for c in (v, p, u, z, b)]) as vpuzb_wrapper:
    for vpuzb_wrapper_local, component in zip(vpuzb_wrapper, (v, p, u, z, b)):
        with component.vector.localForm() as component_local:
            component_local[:] = vpuzb_wrapper_local

J_controlled = mesh.comm.allreduce(dolfinx.fem.assemble_scalar(J_cpp), op=mpi4py.MPI.SUM)
print("Optimal J =", J_controlled)
assert np.isclose(J_controlled, 0.0052941)

Optimal J = 0.005294106310254099

viskex.dolfinx.plot_vector_field(v, "state velocity", glyph_factor=1e-1)

viskex.dolfinx.plot_scalar_field(p, "state pressure")

viskex.dolfinx.plot_vector_field(u, "control", glyph_factor=1e-3)

viskex.dolfinx.plot_vector_field(z, "adjoint velocity", glyph_factor=1e2)

viskex.dolfinx.plot_scalar_field(b, "adjoint pressure")